import cv2

import os
import platform
import sys
from pathlib import Path
from time import sleep
import time

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
from sklearn.cluster import DBSCAN
from numpy import genfromtxt
from numpy import transpose
from numpy import unique
from numpy import where
from PIL import ImageFont, ImageDraw, Image
from collections import Counter
import numpy
import datetime

rgb = [(0, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 0), (255, 255, 255), (100, 100, 255), (255, 100, 100)]


# 获取当前时间
def get_now_datetime():
    now = datetime.datetime.now()
    return str(now.strftime("%Y-%m-%d %H:%M:%S"))


def pattern(img, clusters, yhat, d, img_width, img_height):
    # 控制线的粗细
    pixel = 3
    for cluster in clusters:
        if cluster == -1:
            continue
        rgb_index = cluster + 1
        row_ix = where(yhat == cluster)
        for row_ix_child in row_ix:
            for xy_index in row_ix_child:
                x_centre = d[xy_index, 0]
                y_centre = d[xy_index, 1]
                i = 0
                while i < 5:
                    img = cv2.rectangle(img,
                                        (int(x_centre * img_width) + i, int(y_centre * img_height) + i),
                                        (int(x_centre * img_width) + i, int(y_centre * img_height) + i),
                                        rgb[rgb_index], pixel)
                    img = cv2.rectangle(img,
                                        (int(x_centre * img_width) - i, int(y_centre * img_height) - i),
                                        (int(x_centre * img_width) - i, int(y_centre * img_height) - i),
                                        rgb[rgb_index], pixel)
                    img = cv2.rectangle(img,
                                        (int(x_centre * img_width) + i, int(y_centre * img_height) - i),
                                        (int(x_centre * img_width) + i, int(y_centre * img_height) - i),
                                        rgb[rgb_index], pixel)
                    img = cv2.rectangle(img,
                                        (int(x_centre * img_width) - i, int(y_centre * img_height) + i),
                                        (int(x_centre * img_width) - i, int(y_centre * img_height) + i),
                                        rgb[rgb_index], pixel)
                    i += 1
    return img


def drawing(path, img):
    if not os.path.exists(path):
        return img
    d0 = genfromtxt(path, delimiter=' ')
    d = transpose([d0[:, 1], d0[:, 2]])
    dbscan = DBSCAN(eps=0.1, min_samples=20).fit(d)
    yhat = dbscan.labels_
    clusters = unique(yhat)
    img_width = img.shape[1]
    img_height = img.shape[0]
    number = 0
    count = Counter(dbscan.labels_)

    img = pattern(img, clusters, yhat, d, img_width, img_height)

    fontpath = "simsun.ttc"
    font = ImageFont.truetype(fontpath, 48)
    img_pil = Image.fromarray(img)
    draw = ImageDraw.Draw(img_pil)
    for key in count.keys():
        if key == -1:
            continue
        number = count[key] + number
        draw.text((img_width / 5 * 3, img_height - int(key) * 48 - 400),
                  "第" + str(int(key) + 1) + "组：" + str(count[key]), font=font,
                  fill=rgb[int(key) + 1])
    draw.text((img_width / 5 * 3, img_height - 300), "识别铝棒数量：" + str(number), font=font, fill=(0, 255, 0))
    draw.text((img_width / 5 * 3, img_height - 200), "更新时间：" + str(get_now_datetime()), font=font, fill=(0, 255, 0))
    bk_img = numpy.array(img_pil)

    return bk_img


cap = cv2.VideoCapture(
    'rtsp://admin:lyf188741@192.168.3.20:554/Streaming/Channels/101?transportmode=multicast')

while cap.isOpened():
    c = 0
    while c < 25:
        c += 1
        cap.grab()
    time_start = time.time()
    ret, frame = cap.read()
    frame = drawing("./runs/detect/exp11/labels/K27643130_1_1_0.txt", frame)
    cv2.namedWindow('image', cv2.WINDOW_NORMAL)
    cv2.imshow('image', frame)
    k = cv2.waitKey(20)

    sleep(0.038)
    sleep(0.04)
    time_end_down = time.time()
    print(time_end_down - time_start)

    # q键退出
    if (k & 0xff == ord('q')):
        break

cap.release()
cv2.destroyAllWindows()
